import torch
import torch_npu
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
from lerobot.common.policies.factory import make_policy
from lerobot.configs.policies import PreTrainedConfig


_= torch_npu

def main():
    device = "npu"
    dataset_repo_id = "/home/ascend/dataset/koch_test_dataset"
    ckpt_torch_dir = "/home/ascend/models"

    dataset = LeRobotDataset(dataset_repo_id, episodes=[0])
    

    dataloader = torch.utils.data.DataLoader(
        dataset,
        num_workers=0,
        batch_size=1,
    )
    

    batch = next(iter(dataloader))

    # To device
    for k in batch:
        if isinstance(batch[k], torch.Tensor):
            batch[k] = batch[k].to(device=device, dtype=torch.float32)

    cfg = PreTrainedConfig.from_pretrained(ckpt_torch_dir)
    cfg.pretrained_path = ckpt_torch_dir

    policy = make_policy(cfg, device, ds_meta=dataset.meta)

    warmup_iters = 10
    benchmark_iters = 30

    actions_list = []
    # Warmup
    for _ in range(warmup_iters):
        torch.npu.synchronize()
        action = policy.select_action(batch)
        # 保存单次推理得到的50组关节角度序列到列表中
        # 转为 NumPy 数组 并存储在列表中
        actions_list.append(action.cpu().numpy())  
        policy.reset()
    
        torch.npu.synchronize()
    
    print(actions_list[0])
    # Benchmark
    start_event = torch.npu.Event(enable_timing=True)
    end_event = torch.npu.Event(enable_timing=True)

    start_event.record()
    for _ in range(benchmark_iters):
        action = policy.select_action(batch)
        policy.reset()
    end_event.record()

    # Synchronize and measure time
    torch.npu.synchronize()
    elapsed_time_ms = start_event.elapsed_time(end_event)
    avg_time_per_iter = elapsed_time_ms / benchmark_iters
    print(f"Average execution time per iteration: {avg_time_per_iter:.3f} ms")


if __name__ == "__main__":
    with torch.inference_mode():
        main()
